"""
Includes junction detection evaluation metrics
"""
import numpy as np


# Container for accumulated junction detection results
class JuncResult(object):
    def __init__(self):
        # Record both the original and simplified data
        self.data = {
            "tp": [],
            "fp": [],
            "prob": [],
            "num_gt": 0
        }
        self.data_simple = {
            "tp": [],
            "fp": [],
            "prob": [],
            "num_gt": 0
        }

    # Accumulate the results.
    def update(self, junc_res):
        # Update the normal data
        self.data["tp"].append(junc_res["tp"])
        self.data["fp"].append(junc_res["fp"])
        self.data["prob"].append(junc_res["prob"])
        self.data["num_gt"] += junc_res["num_gt"]

        # Update the simplified data
        self.data_simple["tp"].append(junc_res["tp_simple"])
        self.data_simple["fp"].append(junc_res["fp_simple"])
        self.data_simple["prob"].append(junc_res["prob"])
        self.data_simple["num_gt"] += junc_res["num_gt_simple"]

    # Compute the precision and recall
    def compute_pr(self, mode="original"):
        if mode == "original":
            data = self.data
        elif mode == "simple":
            data = self.data_simple
        else:
            raise ValueError("[Error] Unknown mode!!")

        tp = np.concatenate(data["tp"])
        fp = np.concatenate(data["fp"])
        prob = np.concatenate(data["prob"])


        # Sort in descending order of confidence
        sort_idx = np.argsort(prob)[::-1]
        tp = tp[sort_idx]
        fp = fp[sort_idx]
        prob = prob[sort_idx]

        # Get the cumulative results
        tp_cum = np.cumsum(tp)
        fp_cum = np.cumsum(fp)
        precision = self.div0(tp_cum, tp_cum + fp_cum)
        recall = self.div0(tp_cum, data["num_gt"])

        # Add entries for plotting
        recall = np.concatenate([[0], recall, [1]])
        precision = np.concatenate([[0], precision, [0]])
        # 累积最大精度数组
        precision = np.maximum.accumulate(precision[::-1])[::-1]

        # Compute mAP 实际上是AP（Average Precision）值，即 PR曲线下的面积
        mAP = self.compute_mAP(precision, recall)

        return {
            "precision": precision,
            "recall": recall,
            "prob": prob,
            "mAP": mAP
        }

    # Compute mAP
    @staticmethod
    def compute_mAP(precision, recall):
        return np.sum(precision[1:] * (recall[1:] - recall[:-1]))

    @staticmethod
    def div0(a, b):
        with np.errstate(divide='ignore', invalid='ignore'):
            c = np.true_divide(a, b)
            idx = ~np.isfinite(c)
            c[idx] = np.where(a[idx] == 0, 1, 0)  # -inf inf NaN
        return c


# 将每一个置信度大于阈值的预测点与每一个gt点计算二范数，值小于阈值则认为是匹配点，计算真阳性和假阳性
# 返回 prob（预测点置信度数组）， tp（真阳性布尔数组），fp（假阳性布尔数组），num_gt（gt点数量），
# tp_simple（多个预测为一个gt点的情况），fp_simple（多个预测为一个gt点的情况），
# num_gt_simple（gt被预测到的数量）
# Compute tp fp of junction detection
# TODO: extend to batch size > 1
def compute_tp_fp(junc_pred, junc_gt, dist_thresh=2):
    # junc_pred H W
    # junc_gt H W C(1)
    # dist_thresh = 4
    # Process gt data
    gt_loc = np.where(junc_gt)  # gt_loc是一个大小为 3 的元组，每一项大小为 n ，分别保存 H W C 的索引
    # since 0 is the batch index...

    # n x 2
    gt_loc = np.concatenate([gt_loc[0][..., None], gt_loc[1][..., None]], axis=-1)
    # n
    num_gt = gt_loc.shape[0]

    # Process the prediction data
    zero_thresh = 1e-4

    # 大小为 2 的元组,每一项为 m，分别保存  H W 的索引
    mask = np.where(junc_pred > zero_thresh)
    # Probability predictions
    junc_prob = junc_pred[mask]     # m
    # 0 1 predictions
    # m x 2(H W)
    pred = np.array(mask).T

    # Sorting according to probability
    argsort_idx = np.argsort(junc_prob)[::-1]   # 置信度值从大到小排序的索引
    junc_prob = junc_prob[argsort_idx]
    pred = pred[argsort_idx]
    # Compute the distance differences
    # m x 1 x 2 - 1 x n x 2 = m x n x 2
    diff = np.expand_dims(pred, axis=1) - gt_loc[None, ...]
    # diff 最后一个轴上计算每个元素的二范数
    # m x n
    dist = np.linalg.norm(diff, axis=-1)
    # 若 dist 中值小于 dist_thresh 则为 True， 否则为 False
    # m x n 的一个布尔型数组
    matches = np.less_equal(dist, dist_thresh)

    # Compute tp and fp
    tp = []
    matched = np.zeros(len(gt_loc)) # n
    for m in matches:
        # 检测 m 中是否至少存在一个为 True 的元素
        correct = np.any(m)
        if correct:
            # 返回第一个为 True 的索引
            gt_idx = np.argmax(m)
            tp.append(not matched[gt_idx])
            matched[gt_idx] = 1
        else:
            tp.append(False)
    tp = np.array(tp, bool) # m
    fp = np.logical_not(tp) # m

    # TODO: If we allow multiple predictions to be considered as positive.
    # 确定每个预测至少包含一个 True
    # 大小为 m 的布尔型数组
    tp_simple = np.any(matches, axis=1)  # keeps multiple matches for the same gt point
    # gt被预测的数量
    num_gt_simple = np.sum(np.minimum(np.sum(matches, axis=0), 1))  # buggy
    fp_simple = np.logical_not(tp_simple)

    return {
        "prob": junc_prob,
        "tp": tp,
        "fp": fp,
        "num_gt": num_gt,
        "tp_simple": tp_simple,
        "fp_simple": fp_simple,
        "num_gt_simple": num_gt_simple
    }